import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import spearmanr, pearsonr

import numpy as np
from KunQuant.runner import KunRunner as kr

# def run_kunquant_model(module, inputs: dict, start_time=0, length=None, outputs=None):
#     if length is None:
#         length = next(iter(inputs.values())).shape[0] - start_time

#     N = next(iter(inputs.values())).shape[1]
#     outbuffers = {
#         name: np.zeros((length, N), dtype=np.float32)
#         for name in (outputs or module.getOutputNames())
#     }

#     kr.runGraph("seq", module, inputs, start_time, length, outbuffers)
#     return {name: data.T for name, data in outbuffers.items()}  # (N, D)

# # 使用
# def gen_stock_data2(low, high, stocks, num_time, stddev, dtype="float32"):
#     xopen = np.random.uniform(low, high, size = (stocks, 1)).astype(dtype)
#     # xvol = np.random.uniform(5, 5.2, size = (stocks, 1)).astype(dtype)

#     chopen = np.random.normal(1, stddev, size = (stocks, num_time)).astype(dtype)
#     chopen = np.cumprod(chopen, axis=1, dtype=dtype)
#     outopen = xopen * chopen

#     chopen = np.random.uniform(0.95, 1.05, size = (stocks, num_time)).astype(dtype)
#     outclose = outopen * chopen

#     chopen = np.random.uniform(0.995, 1.12, size = (stocks, num_time)).astype(dtype)
#     outhigh = outopen * chopen
#     outhigh = np.maximum.reduce([outopen, outhigh, outclose])

#     chopen = np.random.uniform(0.9, 1.005, size = (stocks, num_time)).astype(dtype)
#     outlow = outopen * chopen
#     outlow = np.minimum.reduce([outopen, outhigh, outclose, outlow])
    
#     # chopen = np.random.normal(1, stddev, size = (stocks, num_time)).astype(dtype)
#     # chopen = np.cumprod(chopen, axis=1, dtype=dtype)
#     outvol = np.random.uniform(5, 10, size = (stocks, num_time)).astype(dtype)

#     outamount = outvol * outopen * np.random.uniform(0.99, 1.01, size = (stocks, num_time)).astype(dtype)
#     return outopen, outclose, outhigh, outlow, outvol, outamount

# def gen_stock_data(minVal, maxVal, stocksNum, num_time, stddev):
#     xopen = np.random.uniform(minVal, maxVal, size = stocksNum).astype("float32")
#     xvol = np.random.uniform(5, 10, size = stocksNum).astype("float32")
#     outopen = np.empty((stocksNum, num_time), dtype="float32")
#     outclose = np.empty((stocksNum, num_time), dtype="float32")
#     outhigh = np.empty((stocksNum, num_time), dtype="float32")
#     outlow = np.empty((stocksNum, num_time), dtype="float32")
#     outvol = np.empty((stocksNum, num_time), dtype="float32")
#     outamount = np.empty((stocksNum, num_time), dtype="float32")
#     for i in range(num_time):
#         xopen *= np.random.normal(1, stddev, size = (stocksNum)).astype("float32")
#         xclose = xopen * np.random.uniform(0.99, 1.01, size= stocksNum)
#         xhigh = np.maximum.reduce([xopen * np.random.uniform(0.99, 1.03, size= stocksNum), xopen, xclose])
#         xlow = np.minimum.reduce([xopen * np.random.uniform(0.97, 1.01), xopen, xclose, xhigh])
#         xvol *= np.random.normal(1, stddev, size = (stocksNum)).astype("float32")
#         xamount = xvol * xopen * np.random.uniform(0.99, 1.11, size= stocksNum)
#         outopen[:,i] = xopen
#         outclose[:,i] = xclose
#         outhigh[:,i] = xhigh
#         outlow[:,i] = xlow
#         outvol[:,i] = xvol
#         outamount[:,i] = xamount
#     return outopen, outclose, outhigh, outlow, outvol, outamount

# def ST_TS(data: np.ndarray) -> np.ndarray:
#     return np.ascontiguousarray(data.transpose())

# def TS_STs(data: np.ndarray, blocking_num) -> np.ndarray:
#     if data.shape[0] % blocking_num == 0:
#         return np.ascontiguousarray(data.reshape((-1, blocking_num, data.shape[1])).transpose((0, 2, 1)))
#     buf = np.zeros(((data.shape[0] // blocking_num + 1) * blocking_num, data.shape[1]), dtype=data.dtype)
#     buf[:data.shape[0], :] = data
#     return buf

def cal_factor(): 
    # 1. 设置参数
    N = 1000  # 股票数量
    D = 252   # 交易日数量

    # 2. 模拟行情数据（注意：生成 (D, N) 形状）
    np.random.seed(42)
    # def make_array():
    #     return np.random.rand(D, N).astype(np.float32)
    
    # dopen   = np.ascontiguousarray(make_array() * 10 + 10)
    # dhigh   = np.ascontiguousarray(make_array() * 15 + 12)
    # dlow    = np.ascontiguousarray(make_array() * 10 + 8)
    # dclose  = np.ascontiguousarray(make_array() * 10 + 10)
    # dvolume = np.ascontiguousarray(make_array() * 1e8)
    # damount = np.ascontiguousarray(dvolume * dclose)

    # 3. 加载模型
    lib = kr.Library.load(os.path.join("./build/test", "test.dll"))
    modu = lib.getModule("alpha_101")
    blocking_len=modu.blocking_len      #blocking = 0 if is64 else modu.blocking_len

    dopen, dclose, dhigh, dlow, dvol, damount = gen_stock_data2(0.5, 100, D, N, 0.03 if D > 1000 else 0.05)

    # --- 2. 严格验证每个输入 ---
    for name, arr in [("open", dopen), ("high", dhigh), ("low", dlow),
                      ("close", dclose), ("volume", dvol), ("amount", damount)]:
        print(f"{name}: shape={arr.shape}, dtype={arr.dtype}, C={arr.flags['C_CONTIGUOUS']}, F={arr.flags['F_CONTIGUOUS']}")
        if arr.shape != (D, N):
            raise ValueError(f"{name} shape must be ({D}, {N}), got {arr.shape}")
        if arr.dtype != np.float32:
            raise ValueError(f"{name} must be float32, got {arr.dtype}")
        if not arr.flags['C_CONTIGUOUS']:
            raise ValueError(f"{name} is not C-contiguous! Use np.ascontiguousarray()")
  

    # 4. 构建输入字典
    my_input = {"high": TS_STs(dhigh,blocking_len), "low": TS_STs(dlow,blocking_len), "close": TS_STs(dclose,blocking_len), "open": TS_STs(dopen,blocking_len), "volume": TS_STs(dvol,blocking_len), "amount": TS_STs(damount,blocking_len)}    
    # my_input = {
    #     "open": dopen,
    #     "high": dhigh,
    #     "low": dlow,
    #     "close": dclose,
    #     "volume": dvol,
    #     "amount": damount
    # }

    # 5. 准备输出 buffer
    output_names = modu.getOutputNames()
    # print("Model output names:", output_names)

    outbuffers = {}
    for name in output_names:
        # 输出 buffer shape 应为 (length, N)
        outbuffers[name] = np.zeros((D, N), dtype=np.float32)  # 必须 float32

    # 6. 执行计算
    start_time = 0
    length = D
# 创建 Executor 对象
    executor = kr.createSingleThreadExecutor()  # 或 kr.createParallelExecutor()

    outputs_dict = kr.runGraph(
        exec=executor,           # 必须是 Executor 对象
        mod=modu,
        inputs=my_input,
        cur_time=start_time,
        length=length,
        outputs={},              # 可选，一般传空字典
        skip_check=False,
        num_stocks=N
    )

    # 7. 转为 (N, D) 格式返回
    factors = {name: data.T for name, data in outbuffers.items()}  # (N, D)

    print("Factor calculation completed.")
    print("Number of factors:", len(factors))
    
    # 调试：打印第一个因子的部分数值
    first_factor_name = list(factors.keys())[0]
    sample_values = factors[first_factor_name][:5, :5]  # 前5只股票 × 前5天
    print(f"Sample of {first_factor_name}:\n", sample_values)

    # 检查是否全为 0
    if np.allclose(sample_values, 0):
        print("⚠️  Warning: Factor values are all zero! Check input shape, dtype, or model logic.")

    return factors

def factor_analysis(factor_dict: dict, returns: pd.DataFrame, forward_window=5):
    """
    factor_dict: {name: np.array (T, N)} or {name: pd.DataFrame(T, N)}
    returns: pd.DataFrame(T, N), 下一期收益
    """
    results = {}
    
    for name, factor in factor_dict.items():
        if isinstance(factor, np.ndarray):
            factor = pd.DataFrame(factor.T)  # ST -> TS
        
        # Align time index
        common_idx = factor.index.intersection(returns.index)
        fac = factor.loc[common_idx]
        ret = returns.loc[common_idx]

        # Shift factor to predict future return
        fac_lagged = fac.shift(1).dropna()
        ret_forward = ret.rolling(forward_window).sum().shift(-forward_window).loc[fac_lagged.index]

        # Drop NaN
        valid_mask = ~(fac_lagged.isna() | ret_forward.isna())
        fac_clean = fac_lagged[valid_mask].stack()
        ret_clean = ret_forward[valid_mask].stack()

        # IC and Rank IC
        ic = pearsonr(fac_clean, ret_clean)[0]
        rank_ic = spearmanr(fac_clean, ret_clean)[0]
        ir = rank_ic / fac_clean.corr(ret_clean, method='spearman').std() if len(fac_clean) > 1 else np.nan

        turnover = (fac_lagged.diff(1).abs().mean(axis=1)).mean()

        results[name] = {
            'IC': ic,
            'Rank_IC': rank_ic,
            'IR': ir,
            'Turnover': turnover,
            'Mean': fac_clean.mean(),
            'Std': fac_clean.std(),
            'AutoCorr_1D': fac_lagged.corrwith(fac_lagged.shift(1), axis=1).mean()
        }

        # Optional: plot top-bottom spread return
        if False:  # 可视化开关
            deciles = fac_lagged.quantile([0.1, 0.9], axis=1).T
            top = (fac_lagged >= deciles[0.9]).mul(ret_forward, fill_value=0).mean(axis=1)
            bot = (fac_lagged <= deciles[0.1]).mul(ret_forward, fill_value=0).mean(axis=1)
            (top - bot).cumsum().plot(title=f"{name} Long-Short Return")
            plt.show()

    summary = pd.DataFrame(results).T
    summary.sort_values("Rank_IC", key=abs, ascending=False, inplace=True)
    return summary

def combine_factors_ic_weighted(factor_dict, returns, lookback=20):
    from sklearn.preprocessing import RobustScaler
    
    signals = {}
    for t in range(lookback, len(returns)):
        window_factors = {k: v[t-lookback:t] for k, v in factor_dict.items()}
        analysis = factor_analysis(window_factors, returns[t-lookback:t])
        
        weights = analysis['Rank_IC'].fillna(0).clip(-0.3, 0.3)  # 限制极端权重
        weights = weights / weights.abs().sum()  # 归一化
        
        current_factors = pd.DataFrame({k: v[t] for k, v in factor_dict.items()})
        signal = (current_factors * weights).sum(axis=1)
        signals[t] = RobustScaler().fit_transform(signal.values.reshape(-1, 1)).flatten()
    
    return pd.DataFrame(signals).T

def backtest(signal_df: pd.DataFrame, price_df: pd.DataFrame, 
             rebalance_freq='5D', cost_rate=0.001):
    """
    signal_df: (T, N) 综合信号
    price_df: (T, N) 收盘价
    """
    # 对齐时间
    common_idx = signal_df.index.intersection(price_df.index)
    signal = signal_df.loc[common_idx].fillna(0)
    price = price_df.loc[common_idx]

    # 生成持仓（每 5 天调仓）
    position = signal.rolling(2).mean().shift(1)  # 平滑 & lag
    position = position.resample(rebalance_freq).ffill()

    # 收益计算
    ret = price.pct_change().reindex(position.index)
    portfolio_ret = (position * ret).mean(axis=1) - abs(position.diff()).mean(axis=1) * cost_rate

    # 累计净值
    nav = (1 + portfolio_ret).cumprod()
    
    # 绩效指标
    ann_ret = nav.pct_change(periods=252).mean()
    ann_vol = portfolio_ret.std() * np.sqrt(252)
    sharpe = ann_ret / ann_vol
    mdd = (nav / nav.cummax() - 1).min()

    print(f"Sharpe: {sharpe:.2f}, Ann Ret: {ann_ret*100:.2f}%, MDD: {mdd*100:.2f}%")
    
    nav.plot(title="Portfolio NAV")
    plt.show()

    return portfolio_ret, nav

if __name__ == "__main__":
    ret = cal_factor()
    # print(ret)